// Copyright (c) 2021 by Rockchip Electronics Co., Ltd. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "kbt_cam/postprocess.h"

#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>

#include <set>
#include <vector>
#include <algorithm>
typedef struct _BBox {
    float x, y, w, h;
    float score;
    int class_id;
} BBox;

static char* labels[OBJ_CLASS_NUM] = {
    "person",//0
    "bus",//1
    "bicycle",//2
    "car",//3
    "motorbike",//4
    "uav",//5
    "tank",//6
    "soldier",//7
    "a-car",//8
    "plane",//9
    "copter",//10
    "flighter-plane",//11
    "rocket",//12
    "rocket",//13
    "ship",//14
    "ball",//15
    "gun",//16
    "windows",//17
    "xunfeidan"//18
};

inline static int clamp(float val, int min, int max) { return val > min ? (val < max ? val : max) : min; }

inline float deqnt_affine_to_f32(int8_t q, int32_t zp, float scale) {
    return ((float)q - (float)zp) * scale;
}

inline float deq(int8_t val, float scale, int32_t zp) {
    return (static_cast<int>(val) - zp) * scale;
}

static int8_t qnt_f32_to_affine(float f32, int32_t zp, float scale)
{
  float dst_val = (f32 / scale) + zp;
//   int8_t res = (int8_t)__clip(dst_val, -128, 127);
  return dst_val;
}

static float sigmoid(float x) {
    return 1.0f / (1.0f + expf(-x));
}

static void xywh2xyxy(BBox& b) {
    float cx = b.x, cy = b.y;
    b.x = cx - b.w / 2.0f;
    b.y = cy - b.h / 2.0f;
    b.w = cx + b.w / 2.0f;
    b.h = cy + b.h / 2.0f;
}

static int quick_sort_indice_inverse(std::vector<float> &input, int left, int right, std::vector<int> &indices)
{
  float key;
  int key_index;
  int low = left;
  int high = right;
  if (left < right)
  {
    key_index = indices[left];
    key = input[left];
    while (low < high)
    {
      while (low < high && input[high] <= key)
      {
        high--;
      }
      input[low] = input[high];
      indices[low] = indices[high];
      while (low < high && input[low] >= key)
      {
        low++;
      }
      input[high] = input[low];
      indices[high] = indices[low];
    }
    input[low] = key;
    indices[low] = key_index;
    quick_sort_indice_inverse(input, left, low - 1, indices);
    quick_sort_indice_inverse(input, low + 1, right, indices);
  }
  return low;
}

static float iou(const BBox& a, const BBox& b) {
    float xx1 = std::max(a.x, b.x);
    float yy1 = std::max(a.y, b.y);
    float xx2 = std::min(a.w, b.w);
    float yy2 = std::min(a.h, b.h);

    float inter_w = std::max(0.0f, xx2 - xx1);
    float inter_h = std::max(0.0f, yy2 - yy1);
    float inter_area = inter_w * inter_h;

    float area_a = (a.w - a.x) * (a.h - a.y);
    float area_b = (b.w - b.x) * (b.h - b.y);

    return inter_area / (area_a + area_b - inter_area);
}

static float CalculateOverlap(float xmin0, float ymin0, float xmax0, float ymax0, float xmin1, float ymin1, float xmax1,
                              float ymax1)
{
  float w = fmax(0.f, fmin(xmax0, xmax1) - fmax(xmin0, xmin1) + 1.0);
  float h = fmax(0.f, fmin(ymax0, ymax1) - fmax(ymin0, ymin1) + 1.0);
  float i = w * h;
  float u = (xmax0 - xmin0 + 1.0) * (ymax0 - ymin0 + 1.0) + (xmax1 - xmin1 + 1.0) * (ymax1 - ymin1 + 1.0) - i;
  return u <= 0.f ? 0.f : (i / u);
}

// static std::vector<BBox> nms(const std::vector<BBox>& boxes, float threshold) {
//     std::vector<BBox> result;
//     std::vector<BBox> sorted = boxes;
//     std::sort(sorted.begin(), sorted.end(), [](const BBox& a, const BBox& b) {
//         return a.score > b.score;
//     });

//     std::vector<bool> suppressed(sorted.size(), false);
//     for (size_t i = 0; i < sorted.size(); ++i) {
//         if (suppressed[i]) continue;
//         result.push_back(sorted[i]);
//         for (size_t j = i + 1; j < sorted.size(); ++j) {
//             if (iou(sorted[i], sorted[j]) > threshold) {
//                 suppressed[j] = true;
//             }
//         }
//     }
//     return result;
// }
static int nms(int validCount, std::vector<float> &outputLocations, std::vector<int> classIds, std::vector<int> &order,
               int filterId, float threshold)
{
  for (int i = 0; i < validCount; ++i)
  {
    if (order[i] == -1 || classIds[i] != filterId)
    {
      continue;
    }
    int n = order[i];
    for (int j = i + 1; j < validCount; ++j)
    {
      int m = order[j];
      if (m == -1 || classIds[i] != filterId)
      {
        continue;
      }
      float xmin0 = outputLocations[n * 4 + 0];
      float ymin0 = outputLocations[n * 4 + 1];
      float xmax0 = outputLocations[n * 4 + 0] + outputLocations[n * 4 + 2];
      float ymax0 = outputLocations[n * 4 + 1] + outputLocations[n * 4 + 3];

      float xmin1 = outputLocations[m * 4 + 0];
      float ymin1 = outputLocations[m * 4 + 1];
      float xmax1 = outputLocations[m * 4 + 0] + outputLocations[m * 4 + 2];
      float ymax1 = outputLocations[m * 4 + 1] + outputLocations[m * 4 + 3];

      float iou = CalculateOverlap(xmin0, ymin0, xmax0, ymax0, xmin1, ymin1, xmax1, ymax1);

      if (iou > threshold)
      {
        order[j] = -1;
      }
    }
  }
  return 0;
}


static int decode_output(int8_t* input, int grid_h, int grid_w, int stride,
                          std::vector<std::pair<float, float>> anchors,
                          int anchor_offset, float conf_thresh,
                          int32_t zps,
                          float scales,
                          std::vector<float> &boxes, std::vector<float> &objProbs, std::vector<int> &classId)
{
    int grid_len = grid_h * grid_w;
    int num_classes = OBJ_CLASS_NUM;
    int prop_size = 5 + num_classes;
    int num_channels = 3 * prop_size;
    // std::cout<< zps <<"     "<<scales <<std::endl;
    int8_t thres_i8 = qnt_f32_to_affine(conf_thresh,zps,scales);
    int validCount = 0;
    
    // std::vector<float> deq_sigmoid_output;
    // deq_sigmoid_output.resize(num_channels * grid_len);

    // static bool print = false;

    // if (scales.size() < num_channels || zps.size() < num_channels) {
    //     std::cerr << "[Error] scales/zps 尺寸不足：scales.size()=" << scales.size()
    //               << ", zps.size()=" << zps.size()
    //               << ", 应该是 " << num_channels << std::endl;
    //     return;
    // }

    // 打印前10通道的信息
    // for (int ch = 0; ch < std::min(10, num_channels); ++ch) {
    //     int idx = ch * grid_len;
    //     printf("ch=%d raw=%d scale=%.6f zp=%d deq=%.6f sigmoid=%.6f\n",
    //            ch, input[idx], scales[ch], zps[ch],
    //            deq(input[idx], scales[ch], zps[ch]),
    //            sigmoid(deq(input[idx], scales[ch], zps[ch])));
    // }

    // std::cout << "scales.size() = " << scales.size()
    //           << ", zps.size() = " << zps.size()
    //           << ", expected = " << num_channels << std::endl;

    // for (int ch = 0; ch < 10; ++ch) {
    //     int idx = ch * grid_len;  // 每个通道数据起始位置
    //     printf("ch=%d raw=%d scale=%.6f zp=%d deq=%.6f sigmoid=%.6f\n",
    //         ch, input[idx], scales[ch], zps[ch],
    //         deq(input[idx], scales[ch], zps[ch]),
    //         sigmoid(deq(input[idx], scales[ch], zps[ch])));
    // }
    // std::cout << "scales.size() = " << scales.size()
    //       << ", zps.size() = " << zps.size()
    //       << ", expected = " << 3 * (5 + OBJ_CLASS_NUM) << std::endl;


     for (int a = 0; a < 3; ++a) {
        float anchor_w = anchors[anchor_offset + a].first;
        float anchor_h = anchors[anchor_offset + a].second;

        for (int i = 0; i < grid_h; ++i) {
            for (int j = 0; j < grid_w; ++j) {
            
                int pos = i * grid_w + j;
                int base_ch = a * prop_size;

                // 防止越界
                if (base_ch + 4 >= num_channels) continue;

                // if(print == false){
                //     print=true;
                //     for(int i = 0 ;i<prop_size*2;i++){
                //         printf("input[%d]:%d\n",i,input[i]);
                //     }
                // }

                // float val = deq(input[(base_ch + 0) * grid_len + pos], scales, zps);
                // deq_sigmoid_output[(base_ch + 0) * grid_len + pos] = val;

                int8_t obj_int = input[(base_ch + 4) * grid_len + pos];
                // float obj = deq(input[(base_ch + 4) * grid_len + pos],
                //                         scales, zps);
                // if(input[(base_ch + 4) * grid_len + pos]<0)
                //     printf("%d\n",input[(base_ch + 4) * grid_len + pos]);
                if (obj_int < thres_i8) continue;
                
                // if(input[(base_ch + 4) * grid_len + pos]!=-128)
                //     printf("%f\n",input[(base_ch + 4) * grid_len + pos]);
                // printf("obj:%f,input:%d\n",obj,input[(base_ch + 4) * grid_len + pos]);

                float bx = deq(input[(base_ch + 0) * grid_len + pos], scales, zps);
                float by = deq(input[(base_ch + 1) * grid_len + pos], scales, zps);
                float bw = deq(input[(base_ch + 2) * grid_len + pos], scales, zps);
                float bh = deq(input[(base_ch + 3) * grid_len + pos], scales, zps);

                float cx = ((bx * 2.0f - 0.5f) + j) * (float)stride;
                float cy = ((by * 2.0f - 0.5f) + i) * (float)stride;
                float w = powf(bw * 2.0f, 2.0f) * (float)anchor_w;
                float h = powf(bh * 2.0f, 2.0f) * (float)anchor_h;

                // int test1 =55;
                // float test_x = deq(test1, scales, zps);
                // float test_cx = ((test_x * 2.0f - 0.5f) + j) * stride;

                // int test2 =56;
                // float test_y = deq(test2, scales, zps);
                // float test_cy = ((test_y * 2.0f - 0.5f) + j) * stride;
                
                // printf("cx:%f,cy:%f\n",test_cx,test_cy);

                int8_t max_score = 0;
                int max_cls = -1;
                for (int c = 0; c < OBJ_CLASS_NUM; ++c) {
                    int cls_ch = base_ch + 5 + c;
                    int cls_idx = cls_ch * grid_len + pos;
                    // int cls_val = input[cls_idx];
                    // float cls_val = deq(input[cls_idx], scales, zps)
                    if (input[cls_idx] > max_score) {
                        max_score = input[cls_idx];
                        max_cls = c;
                    }
                }
                // float final_score = max_score * obj;
                // printf("%f,%f,%f,%f\n",cx, cy, w, h);
                // 改变识别目标，5是uav，15是ball
                if (max_cls != 5)
                    continue;
                if (max_score >= conf_thresh) {
                    // printf("%f,%f,%f,%f\n",cx, cy, w, h);
                    float final_score =  deq(max_score, scales, zps) * deq(obj_int, scales, zps);
                    // out_boxes.push_back({cx, cy, w, h, final_score, max_cls});
                    // printf("%f,%f,%f,%f\n",cx, cy, w, h);
                    objProbs.push_back(final_score);
                    classId.push_back(max_cls);
                    validCount++;
                    boxes.push_back(cx);
                    boxes.push_back(cy);
                    boxes.push_back(w);
                    boxes.push_back(h);
                }
            }
        }
    } 
    return validCount;
}
// 加入调试输出
                // if (cls_ch < num_channels) {
                //     float cls_val = sigmoid(deq(input[cls_idx], scales[cls_ch], zps[cls_ch]));
                //     float score = cls_val * obj;

                //     printf("cls_ch=%d cls_idx=%d raw=%d scale=%.6f zp=%d deq=%.6f sigmoid=%.6f score=%.6f\n",
                //         cls_ch, cls_idx, input[cls_idx], scales[cls_ch], zps[cls_ch],
                //         deq(input[cls_idx], scales[cls_ch], zps[cls_ch]), cls_val, score);

                //     if (score >= conf_thresh) {
                //         out_boxes.push_back({cx, cy, w, h, score, target_class});
                //     }
                // } else {
                //     printf("[Skip] cls_ch=%d 超出通道范围 num_channels=%d\n", cls_ch, num_channels);
                // }
//  if (!saved) {
    //                 FILE* fp = fopen("/tmp/deq_output.bin", "wb");
    //                 if (!fp) {
    //                     std::cerr << "[ERROR] Failed to open file for saving dequantized output." << std::endl;
    //                 } else {
    //                     fwrite(deq_sigmoid_output.data(), sizeof(float), deq_sigmoid_output.size(), fp);
    //                     fclose(fp);
    //                     std::cout << "[INFO] Dequantized output saved to /tmp/deq_output.bin" << std::endl;
    //                     saved = true;  // 只保存一次
    //                 }
    //             }
                // float max_score = 0.0f;
                // int max_cls = -1;
                // for (int c = 0; c < OBJ_CLASS_NUM; ++c) {
                //     float cls = sigmoid(deq(input[idx + 5 + c], scales[5 + c], zps[5 + c]));
                //     if (cls > max_score) {
                //         max_score = cls;
                //         max_cls = c;
                //     }
                // }
                // float final_score = max_score * obj;
                
                // if (max_cls != 15) continue;
                // if (obj < 0.55f) continue;
                // std::cout << "obj = " << obj << ", class = " << max_cls << ", score = " << final_score << std::endl;
                
                
                // if (final_score >= conf_thresh) {
                //     out_boxes.push_back({cx, cy, w, h, final_score, max_cls});
                // }
  


int post_process(int8_t* input0, int8_t* input1, int8_t* input2,
                 int model_in_h, int model_in_w,
                 float conf_thresh, float nms_thresh,
                 BOX_RECT pads, float scale_w, float scale_h,
                 std::vector<int32_t>& zps, std::vector<float>& scales,
                 detect_result_group_t* group)
{
    // std::vector<BBox> raw_boxes;
    std::vector<float> filterBoxes;
    std::vector<float> objProbs;
    std::vector<int> classId;
    std::vector<std::pair<float, float>> anchors = {
        {10,13}, {16,30}, {33,23},
        {30,61}, {62,45}, {59,119},
        {116,90}, {156,198}, {373,326}
    };
    int validCount = 0;
    
    memset(group, 0, sizeof(detect_result_group_t));

    validCount += decode_output(input0, 80, 80, 8, anchors, 0, conf_thresh, zps[0], scales[0], filterBoxes, objProbs, classId);
    // printf("raw_boxes size: %zu\n", raw_boxes.size());
    validCount += decode_output(input1, 40, 40, 16, anchors, 3, conf_thresh, zps[1], scales[1], filterBoxes, objProbs, classId);
    // printf("raw_boxes size: %zu\n", raw_boxes.size());
    validCount += decode_output(input2, 20, 20, 32, anchors, 6, conf_thresh, zps[2], scales[2], filterBoxes, objProbs, classId);


    // printf("validCount: %d\n",validCount);
    if(!validCount)
        return 0;
    // printf("validCount: %d\n",validCount);
    // printf("raw_boxes size: %zu\n", raw_boxes.size());
    // for (auto& box : raw_boxes)
    //     xywh2xyxy(box);
    
    std::vector<int> indexArray;
    for (int i = 0; i < validCount; ++i)
        indexArray.push_back(i);
    
    quick_sort_indice_inverse(objProbs, 0, validCount - 1, indexArray);

    std::set<int> class_set(std::begin(classId), std::end(classId));

    for (auto c : class_set)
        nms(validCount, filterBoxes, classId, indexArray, c, nms_thresh);
    
    int last_count = 0;
    group->count = 0;

    //只保留一个框
    float max_conf = -1.0f;
    int max_index = -1;

    for (int i = 0; i < validCount; ++i)
    {
        if (indexArray[i] == -1)
            continue;

        int n = indexArray[i];
        float obj_conf = objProbs[i];

        if (obj_conf > max_conf) {
            max_conf = obj_conf;
            max_index = n;
        }
    }
    //只保留一个框

    if (max_index >= 0)
    {
        float cx = filterBoxes[max_index * 4 + 0];
        float cy = filterBoxes[max_index * 4 + 1];
        float w  = filterBoxes[max_index * 4 + 2];
        float h  = filterBoxes[max_index * 4 + 3];

        float x1 = (cx - w / 2.0f - pads.left);
        float y1 = (cy - h / 2.0f - pads.top);
        float x2 = (cx + w / 2.0f - pads.left);
        float y2 = (cy + h / 2.0f - pads.top);

        int id = classId[max_index];
        char *label = labels[id];

        group->results[0].box.left   = (int)(clamp(x1, 0, model_in_w) / scale_w);
        group->results[0].box.top    = (int)(clamp(y1, 0, model_in_h) / scale_h);
        group->results[0].box.right  = (int)(clamp(x2, 0, model_in_w) / scale_w);
        group->results[0].box.bottom = (int)(clamp(y2, 0, model_in_h) / scale_h);
        group->results[0].prop       = max_conf;
        strncpy(group->results[0].name, label, OBJ_NAME_MAX_SIZE);

        group->count = 1;
    }
    else {
        group->count = 0; // 没有检测框
    }

    return group->count;


    // /* box valid detect target */
    // for (int i = 0; i < validCount; ++i)
    // {
    //     if (indexArray[i] == -1 || last_count >= OBJ_NUMB_MAX_SIZE)
    //         continue;
    //     int n = indexArray[i];

    //     float cx = filterBoxes[n * 4 + 0];
    //     float cy = filterBoxes[n * 4 + 1];
    //     float w  = filterBoxes[n * 4 + 2];
    //     float h  = filterBoxes[n * 4 + 3];

    //     // float x1 = filterBoxes[n * 4 + 0] - pads.left;
    //     // float y1 = filterBoxes[n * 4 + 1] - pads.top;
    //     // float x2 = x1 + filterBoxes[n * 4 + 2] - pads.left;
    //     // float y2 = y1 + filterBoxes[n * 4 + 3] - pads.top;
    //     float x1 = (cx - w / 2.0f - pads.left);
    //     float y1 = (cy - h / 2.0f - pads.top) ;
    //     float x2 = (cx + w / 2.0f - pads.left);
    //     float y2 = (cy + h / 2.0f - pads.top) ;
    //     int id = classId[n];
    //     float obj_conf = objProbs[i];

    //     group->results[last_count].box.left = (int)(clamp(x1, 0, model_in_w) / scale_w);
    //     group->results[last_count].box.top = (int)(clamp(y1, 0, model_in_h) / scale_h);
    //     group->results[last_count].box.right = (int)(clamp(x2, 0, model_in_w) / scale_w);
    //     group->results[last_count].box.bottom = (int)(clamp(y2, 0, model_in_h) / scale_h);
    //     group->results[last_count].prop = obj_conf;
    //     char *label = labels[id];
    //     strncpy(group->results[last_count].name, label, OBJ_NAME_MAX_SIZE);

    //     // printf("result %2d: (%4d, %4d, %4d, %4d), %s\n", i, group->results[last_count].box.left,
    //     // group->results[last_count].box.top,
    //     //        group->results[last_count].box.right, group->results[last_count].box.bottom, label);
    //     last_count++;
    // }
    // group->count = last_count;

    // return last_count;
}





    // auto final_boxes = nms(raw_boxes, nms_thresh);
    // printf("final_boxes size: %zu\n", final_boxes.size());

    // if(group->count = std::min((int)final_boxes.size(), OBJ_NUMB_MAX_SIZE))
    //     return 0;

    // std::cout<<group->count<<std::endl;
    
    // for (int i = 0; i < group->count; ++i) {
    //     const auto& b = final_boxes[i];
    //     int left = std::max(int((b.x - pads.left) * scale_w), 0);
    //     int top = std::max(int((b.y - pads.top) * scale_h), 0);
    //     int right = std::min(int((b.w - pads.left) * scale_w), model_in_w);
    //     int bottom = std::min(int((b.h - pads.top) * scale_h), model_in_h);

    //     group->results[i].box = {left, right, top, bottom};
    //     group->results[i].prop = b.score;
    //     snprintf(group->results[i].name, OBJ_NAME_MAX_SIZE, "%d", b.class_id);
    // }

 
void deinitPostProcess()
{
}
